#ifndef BUILDINGEXTRACTOR_H
#define BUILDINGEXTRACTOR_H

#include <opencv2/opencv.hpp>
#include <onnxruntime_cxx_api.h>
#include <string>
#include <vector>
#include <cmath>

class BuildingExtractor
{
public:
    BuildingExtractor();
    void initialize(const std::string& modelPath);
    cv::Mat extract(const cv::Mat& image);

private:
    cv::Mat preprocess(const cv::Mat& image);
    std::vector<float> softmax(const std::vector<float>& input); 
    cv::Mat labelToRGB(const cv::Mat& mask);

    Ort::Env env{ nullptr };
    Ort::SessionOptions sessionOptions{ nullptr };
    Ort::Session session{ nullptr };
    std::vector<const char*> inputNames;
    std::vector<const char*> outputNames;
    std::vector<int64_t> InputShape;
    std::vector<int64_t> OutputShape;
};

#endif 